package org.sean.framework.canary.config;

import io.github.resilience4j.circuitbreaker.CircuitBreakerConfig;
import io.github.resilience4j.timelimiter.TimeLimiterConfig;
import org.apache.kafka.common.header.Header;
import org.apache.kafka.common.header.Headers;
import org.sean.framework.canary.core.CanaryProperties;
import org.sean.framework.canary.feign.CanaryInfoHolder;
import org.sean.framework.logging.Logger;
import org.sean.framework.util.StringUtil;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.cloud.circuitbreaker.resilience4j.Resilience4JCircuitBreakerFactory;
import org.springframework.cloud.circuitbreaker.resilience4j.Resilience4JConfigBuilder;
import org.springframework.cloud.client.circuitbreaker.Customizer;
import org.springframework.cloud.loadbalancer.annotation.LoadBalancerClients;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Import;
import org.springframework.kafka.core.KafkaTemplate;
import org.springframework.kafka.listener.RecordInterceptor;

import java.time.Duration;
import java.util.HashMap;
import java.util.Map;

/**
 * 启动配置
 * Created by Sean.xie on 2018/11/26.
 */
@Configuration
@Import({CanaryProperties.class})
@LoadBalancerClients(defaultConfiguration = {LoadBalancerConfiguration.class})
public class CanaryAutoConfiguration {

    @Configuration
    @ConditionalOnClass(Resilience4JCircuitBreakerFactory.class)
    public static class CircuitBreakerConfiguration {

        @Bean
        public Customizer<Resilience4JCircuitBreakerFactory> defaultCustomizer() {
            return factory -> factory.configureDefault(id -> new Resilience4JConfigBuilder(id)
                    .timeLimiterConfig(TimeLimiterConfig.custom().timeoutDuration(Duration.ofSeconds(5)).build())
                    .circuitBreakerConfig(CircuitBreakerConfig.ofDefaults())
                    .build());
        }
    }

    @Configuration
    @ConditionalOnClass(KafkaTemplate.class)
    public static class KafkaConfiguration {
        @Bean
        @Autowired
        public RecordInterceptor recordInterceptor() {
            return record -> {
                Headers kHeaders = record.headers();
                if (kHeaders != null) {
                    String traceId = null;
                    String span = null;
                    Map<String, String> headers = new HashMap<>();
                    for (Header header : kHeaders) {
                        String key = header.key();
                        String value = new String(header.value());
                        headers.put(key, value);
                        if (Logger.KEY_TRACE.equalsIgnoreCase(key)) {
                            traceId = value;
                        }
                        if (Logger.KEY_SPAN.equalsIgnoreCase(key)) {
                            span = value;
                        }
                    }
                    if (StringUtil.isNotEmpty(span) && StringUtil.isNotEmpty(traceId)) {
                        CanaryInfoHolder.handle(headers, traceId, span);
                    } else {
                        CanaryInfoHolder.cleanHeaders();
                    }
                }
                return record;
            };
        }
    }
}

